import numpy as np
from PIL import Image
import math


def get_frequency_linear(width, wave_length):
    out = np.ones(width, np.int)
    for i in range(width):
        if i % (2 * wave_length) < wave_length:
            out[i] = -1
    return out

mask_dict = {}
def get_frequency_mask(img_size, wave_length, angle):
    # mask_name = f'{size[0]}-{size[1]}-{wave_length}-{angle}'
    # if mask_name in mask_dict:
    #     return mask_dict[mask_name]

    size = int(math.ceil((img_size[0]**2 + img_size[1]**2)**.5))
    # wave_length = math.ceil(max(img_size) * wave_length / 100)
    freq_mask = np.ones((size, size), np.float32)
    wave_length2 = 2 * int(wave_length + 0.5)
    for i in range(freq_mask.shape[0]):
        if i % wave_length2 < wave_length:
            freq_mask[i, :] = -1
    
    freq_mask = Image.fromarray(freq_mask)
    freq_mask = freq_mask.rotate(angle)
    freq_mask = np.asarray(freq_mask)
    
    x_start = int((size - img_size[0]) / 2 + 0.5)
    y_start = int((size - img_size[1]) / 2 + 0.5)
    freq_mask = freq_mask[x_start: x_start+img_size[0], y_start: y_start+img_size[1]]

    # mask_dict[mask_name] = freq_mask
    return freq_mask


if __name__ == '__main__':
    print(get_frequency_mask((10, 10), 2, 90))